import pandas as pd
from flask import Flask, send_from_directory
import os
import sqlite3
from datetime import datetime

'''
flask 实现文件下载
'''
app = Flask(__name__)


def query_data(sql):
    con = sqlite3.connect('./data/userTest.db')
    try:
        cur = con.cursor()
        cur.execute(sql)

        return cur.description, cur.fetchall()
    finally:
        con.close()


def generate_execl(data_dir: str, file_name: str):
    file_path = os.path.join(data_dir, file_name)

    cols_name, datas = query_data('select * from user;')
    df = pd.DataFrame(
        data=[list(row) for row in datas],
        columns=[col[0] for col in cols_name]
    )
    df.to_excel(f'{file_path}', sheet_name='pvuv', index=False)


@app.route('/download_pvuv_excel')
def download_pvuv_excel():
    print(app.root_path)
    data_dir = os.path.join(app.root_path, 'downloads')
    try:
        os.makedirs(data_dir)
    except FileExistsError as err:
        print(err)
    now_time = datetime.now().strftime("%Y%m%d_%H%M%S")

    file_name = f'pvuv_{now_time}.xlsx'
    generate_execl(data_dir, file_name)

    return send_from_directory(data_dir, file_name, as_attachment=True)


if __name__ == "__main__":
    app.run()